import math
import numpy as np

def Lboundary(x, lx):
    res=x
    while x>=lx:
        x-=lx
    while x<0:
        x+=lx
    return x

def boundary(pos, lx,ly,lz):
    pos[0]=Lboundary(pos[0],lx)
    pos[1]=Lboundary(pos[1],ly)
    pos[2]=Lboundary(pos[2],lz)
    return pos

def anint(x, lx):
    res=x
    if x>0.5*lx:
        res-=lx
    elif x<-0.5*lx:
        res+=lx
    return res

def distancia(lbox,pos,i,j):
    dis2=0
    drx=pos[i][0]-pos[j][0]
    drx=anint(drx, lbox[0])
    dry=pos[i][1]-pos[j][1]
    dry=anint(dry, lbox[1])
    drz=pos[i][2]-pos[j][2]
    #drz=anint(drz, lbox[2])
    dis2=drx*drx+dry*dry+drz*drz
    return  [math.sqrt(dis2),drx,dry,drz]

def EpotenDip(m,diametros,lbox,datos):
    ePot=0.
    for i in range(len(datos)):
        for j in range(len(datos)):
            if i!=j:
                u0=m*m*math.pow(diametros[i],3)*math.pow(diametros[j],3)
                [dis, dx,dy,dz]= distancia(lbox,datos,i,j)
                Co= dz/dis
                Co2=Co*Co
                dis3=dis*dis*dis
                potdm = u0*(1.-3.*Co2)/dis3
                ePot+=potdm
    return 0.5*ePot/len(datos)

def EpotenKlin(m,diameter,lbox,datos):
    kappa=100.0
    ePot=0.
    for i in range(len(datos)):
        for j in range(len(datos)):
            if i!=j:
                sigma = 0.5*(diameter[i]+diameter[j])
                u0=m*m*math.pow(diameter[i],3.0)*math.pow(diameter[j],3.0)
                u0=m*m #IMPORTANTE!!! ESTO HAY QUE QUITARLO
                [dis, dx,dy,dz]= distancia(lbox,datos,i,j)
                dis3=dis*dis*dis
                potdm = 3.*u0/kappa/sigma*math.exp(kappa*(sigma-dis))
                ePot+=potdm
    return 0.5*ePot/len(datos)

def eKin(diametro, vel):
    ek=0.0;
    np=0 
    for i in range(len(vel)):
        mass = math.pow(diametro[i],3.0)
        vx=vel[i][0]
        vy=vel[i][1]
        vz=vel[i][2]
        v2=mass*(vx*vx+vy*vy+vz*vz)
        ek+=v2
    return 0.5*ek/len(vel)

def importConfig(name):
    f = open(name,'r')
    tx = f.readline()
    tx = tx.replace('\n','')
    val = tx.split(' ')
    aux=[]
    for v in val:
        if v!='':
            aux.append(v)
    t=float(aux[1])
    tcorr=float(aux[3])
    icorr=int(aux[5])
    npart=int(aux[7])
    lx=float(aux[9])
    ly=float(aux[10])
    lz=float(aux[11])
    tx = f.readline()
    tx = tx.replace('\n','')
    val = tx.split(' ')
    aux=[]
    for v in val:
        if v!='':
            aux.append(v)
    sigma=[float(aux[1]),float(aux[2]),float(aux[3]),float(aux[4]),float(aux[5]),float(aux[6])]
    pos=[]
    vel=[]
    force=[]
    diameter=[]
    
    for i in range(npart):
        tx = f.readline()
        tx = tx.replace('\n','')
        val = tx.split(' ')
        aux=[]
        for v in val:
            if v!='':
                aux.append(v)
    
        xpos=boundary([float(aux[0]),float(aux[1]),float(aux[2])], lx,ly,lz)
        pos.append(xpos)
        vel.append([float(aux[3]),float(aux[4]),float(aux[5])])
        force.append([float(aux[6]),float(aux[7]),float(aux[8])])
        diameter.append(float(aux[9]))
    return [[lx,ly,lz],pos,vel,force,diameter,sigma]


def forceDipij(diameter, m,lbox, pos,i,j):
    u0=m*m*math.pow(diameter[i],3.0)*math.pow(diameter[j],3.0)
    [dis,dx,dy,dz] = distancia(lbox,pos,i,j)
    Co= dz/dis
    Co2=Co*Co
    fx=-3.*u0*(5.*Co2-1.)*dx/math.pow(dis,5.0)
    fy=-3.*u0*(5.*Co2-1.)*dy/math.pow(dis,5.0)
    fz=-3.*u0*(5.*Co2-3.)*dz/math.pow(dis,5.0)
    return [fx,fy,fz]


def forceDip(diameter, m,lbox, pos):
    forc=np.zeros((len(pos),3))  
    for i in range(len(pos)):
        for j in range(len(pos)):
            if i!=j:
                [fx,fy,fz]=forceDipij(diameter, m,lbox, pos,i,j)            
                forc[i][0]+=fx
                forc[i][1]+=fy
                forc[i][2]+=fz
    return forc

def forceKlingembergij(diameter, m,lbox, pos,i,j):
    kappa=100.0
    sigma = 0.5*(diameter[i]+diameter[j])
    u0=m*m*math.pow(diameter[i],3.0)*math.pow(diameter[j],3.0)
    u0=m*m  #IMPORTANTE!!! ESTO HAY QUE QUITARLO
    [dis,dx,dy,dz] = distancia(lbox,pos,i,j)
    fff = u0 *math.exp(kappa*(sigma-dis))/dis

    fx=fff*dx
    fy=fff*dy
    fz=fff*dz
    return [fx,fy,fz]


def forceKlingemberg(diameter, m,lbox, pos):
    forc=np.zeros((len(pos),3))  
    for i in range(len(pos)):
        for j in range(len(pos)):
            if i!=j:
                [fx,fy,fz]=forceKlingembergij(diameter, m,lbox, pos,i,j)            
                forc[i][0]+=fx
                forc[i][1]+=fy
                forc[i][2]+=fz
    return forc


def stressTensor(diameter, m,lbox,vel,pos,volume):
    npart = len(pos)
    sumv=np.zeros((3,3))
    sumf=np.zeros((3,3))
    sigma=np.zeros((3,3))
    #Parte dependiente de la velocidad
    for i in range(npart):
        mass = math.pow(diameter[i],3.0)
        #mass=1.0
        for i1 in range(3):
            for i2 in range(3):
                sumv[i1][i2]+=mass*vel[i][i1]*vel[i][i2]
    #Parte dependiente de la fuerza
    for i in range(npart-1):
        for j in range(i+1,npart):
            forceD=forceDipij(diameter, m,lbox, pos,i,j)
            forceK=forceKlingembergij(diameter, m,lbox, pos,i,j)
            [dis,dx,dy,dz] = distancia(lbox,pos,i,j)
            dis=[dx,dy,dz]
            for i1 in range(3):
                for i2 in range(3):
                    sumf[i1][i2]+=(forceD[i2]+forceK[i2])*dis[i1]
    #Sumamos ambas partes
    for i1 in range(3):
            for i2 in range(3):
                sigma[i1][i2]=(sumv[i1][i2]+sumf[i1][i2])/volume
    return [sumv,sumf,sigma]

def calcForcesPart(diameter, m, lbox, pos, i):
    force = np.zeros((3))
    for j in range(len(pos)):
        if i!=j:
            fk = forceKlingembergij(diameter, m,lbox, pos,i,j)
            fd = forceDipij(diameter, m,lbox, pos,i,j)
            force += np.array(fk)+np.array(fd)
    return force
            

u0=0*100.0/8.0
m=math.sqrt(u0)
name = 'config.dat'
[lbox,pos,vel,force,diameter,sigmaConfig] = importConfig(name)
volume = lbox[0]*lbox[1]*lbox[2]
epotenDip = EpotenDip(m,diameter,lbox,pos)
ekinet = eKin(diameter,vel)
ekling = EpotenKlin(m,diameter,lbox,pos)
[sumv,sumf,sigma] = stressTensor(diameter, m,lbox,vel,pos,volume)

print "Energia dipolar: ",epotenDip
print "Energia cinetica: ",ekinet
print "E Klingemberg: ",ekling
print "Tensor esfuerzos espresso: ",sigmaConfig
print "Tensor esfuerzos calculado: ",[sigma[0][0], sigma[0][1],sigma[0][2],sigma[1][1],sigma[1][2],sigma[2][2]]


iii=0
forceNeo = calcForcesPart(diameter, m, lbox, pos, iii)
print forceNeo
print force[iii]